# -*- coding: utf-8 -*-
import os
import sys,getopt
import numpy as np
import time
import random
import argparse

MC_FILE_ROOT = os.path.join('..', '..', 'CommonFiles', 'SemanticMidFiles', 'MonteCarloEvaluation')


np.random.seed(0)

randomN = 1000000

def binarySearchforA(A,target):
    i = 0
    j = len(A)-1
    while i < j:
        mid = (j + i) / 2
        mid = int(mid)
        if target >= A[mid]:
            i = mid + 1
        else:
            j = mid
    return mid

def listgen(ruleset,gendict):
    for home, dirs, files in os.walk(ruleset):
        for filename in files:
            category = filename.split('.')[0]
            f = open(os.path.join(home, filename),'r')
            tmp = {}
            line = f.readline()
            while line:
                line = line.strip('\n')
                linestr, lineprob = line.split('\t')
                tmp[linestr] = float(lineprob)
                line = f.readline()

            gendict[category] = sorted(tmp.items(), key=lambda tmp:tmp[1], reverse=True)
            f.close()

def randomChooseN(structfile,gendict,outputfile):
    global randomN
    patternlist = list()
    patternproblist = list()
    patternproblistSum = [0]

    with open(structfile,'r') as f:
        pw = f.readline()
        idx = 0
        while pw:
            pw = pw.strip('\n')
            pwstr = '\t'.join(pw.split('\t')[:-1])
            pwprob = float(pw.split('\t')[-1])
            patternlist.append(pwstr)
            patternproblist.append(pwprob)
            patternproblistSum.append(patternproblistSum[-1]+pwprob)
            idx += 1
            pw = f.readline()

    factorlistDict = dict()
    factorProblistDict = dict()
    factorProblistSumDict = dict()
    usedPwdProb = dict()
    unmatched_cnt = 0
    with open(outputfile, 'w') as f:
        finallist = list()
        i = 0
        while i < randomN:
            randomVal = random.random()
            curidx = binarySearchforA(patternproblistSum,randomVal)

            basefactor = ""
            baseProb = patternproblist[curidx]

            curpattern = patternlist[curidx].split('\t')

            for factor in curpattern:
                factorProblist = list()
                factorProblistSum = [0]
                factorlist = list()
                if factor in factorlistDict:
                    factorlist =  factorlistDict[factor]
                    factorProblist = factorProblistDict[factor]
                    factorProblistSum = factorProblistSumDict[factor]
                else:
                    pass

                    for tpl in gendict[factor]:
                        if len(tpl[0]) > 4 and tpl[0][:4] == 'P_U_':
                            this_sum = int(tpl[0].split('-')[-1])
                            total_prob = this_sum * tpl[1]
                            #print(tpl[0], this_sum, total_prob)
                            factorProblist.append(tpl[1])
                            factorlist.append(tpl[0])
                            factorProblistSum.append(factorProblistSum[-1] + total_prob)
                        else:
                            factorProblist.append(tpl[1])
                            factorlist.append(tpl[0])
                            factorProblistSum.append(factorProblistSum[-1]+tpl[1])

                    factorProblistSumDict[factor] = factorProblistSum
                    factorlistDict[factor] = factorlist
                    factorProblistDict[factor] = factorProblist

                randomVal = random.random()
                curidx = binarySearchforA(factorProblistSum, randomVal)
                curProb = factorProblist[curidx]
                curfactor = factorlist[curidx]

                baseProb *= curProb
                basefactor += curfactor

            if(len(basefactor) >= 6):
                i += 1
                if basefactor in usedPwdProb.keys():
                    if baseProb < usedPwdProb[basefactor]:
                        #print("Find smaller prob pwd:", basefactor, baseProb, usedPwdProb[basefactor])
                        unmatched_cnt += 1
                        usedPwdProb[basefactor] = baseProb
                    elif baseProb > usedPwdProb[basefactor]:
                        #print("Find larger prob pwd", basefactor, baseProb, usedPwdProb[basefactor])
                        unmatched_cnt += 1
                        baseProb = usedPwdProb[basefactor]
                    else:
                        pass
                else:
                    usedPwdProb[basefactor] = baseProb
                finallist.append([basefactor, baseProb])
        
        print("unmatched count:", unmatched_cnt)
        for i in range(len(finallist)):
            if finallist[i][0] in usedPwdProb.keys():
                finallist[i][1] = usedPwdProb[finallist[i][0]]
        finallist.sort(reverse=True, key=lambda finallist:finallist[1])
        for item in finallist:
            #f.write(str(item[0]) + '\t' + str(item[1]) +'\n')
            f.write(str(item[1]) + '\n')

'''
    该文件针对SE#PCFG进行蒙特卡洛模拟中的随机采样
    包含了四个参数，分别为：
    
    1. dataset_name：        输入参数，具体要做模拟的训练库名称
    2. dataset_root：        输入参数，分词结果存放的根目录
    3. ruleset_name：        输入参数，训练库使用哪个版本的分词结果。注意这里要使用Rules文件夹的全称，因为可能会有批量的实验
    4. output_root：         输入参数，蒙特卡洛采样结果存放目录
'''
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_name', type=str, nargs='+', default=['CSDN_100w_20210806', 'GMAIL_100w_20210806', 'EYEEM_100w_20210806', 'FR0_100w_20210806'])
    parser.add_argument('--dataset_root', default=os.path.join('..', 'Results', 'Result_20221029'))
    parser.add_argument('--ruleset_name', default='Rules')
    parser.add_argument('--output_root', default=os.path.join(MC_FILE_ROOT, 'MonteCarlo_20221213', 'mc_models'))

    args = parser.parse_args()

    dataset_name_list = args.dataset_name
    dataset_root = args.dataset_root
    output_root = args.output_root
    ruleset_name = args.ruleset_name

    if not os.path.exists(output_root):
        os.makedirs(output_root)

    for one_dataset in dataset_name_list:
        time_start = time.time()

        ruleset = os.path.join(dataset_root, one_dataset, ruleset_name)
        structfile = os.path.join(dataset_root, one_dataset, '%s%s' % (one_dataset, '_genstruct'))
        outputfile = os.path.join(output_root, '%s%s' % (one_dataset, '_MC_model'))

        gendict = {}
        listgen(ruleset, gendict)
        randomChooseN(structfile, gendict, outputfile)

        time_end = time.time()
        print('totally cost ', time_end - time_start)

if __name__ == '__main__':
    main()

